
from numpy import inf,  zeros, tile, nonzero, maximum, minimum,ceil,floor,argsort,array,argmax,sort,mod,cumsum,sum,arange,linspace,ones
from numpy.random import rand,randint,randn
import sys



class GA(object):
    
    
    @staticmethod
    def default_values(optinfo):
        

        if not 'PROPelite' in optinfo.keys():
            optinfo['PROPelite'] = 0.05
        if not 'PROPxover' in optinfo.keys():
            optinfo['PROPxover'] = 0.6
        if not 'PROPmutation' in optinfo.keys():
            optinfo['PROPmutation'] = 0.3
            
        if not 'FCscale' in optinfo.keys(): 
            optinfo['FCscale'] = 'Ranking'
             
        if not 'FCselection' in optinfo.keys():
            optinfo['FCselection'] = 'StocUnfiform'
        
        if not 'FCxover' in optinfo.keys():
            optinfo['FCxover'] ='Intermediate'
            
        if not 'FCmutation' in optinfo.keys():
            optinfo['FCmutation'] = 'Gaussian'
            #optinfo['FCmutation'] = 'Uniform'
            optinfo['SCALEmutation'] = 1
            optinfo['shrinkMutation'] = 0.5
            
        if not 'Minterval' in optinfo.keys():
            optinfo['Minterval']=10
            
        if not 'PCMigration' in optinfo.keys():   #of the total population, should be smalled than the number of worker
            optinfo['PCMigration']=0.1
        
        if not 'PCParents' in optinfo.keys():   
            optinfo['PCParents']=1
        
        if  optinfo['FCxover'] == 'LinearComb':
            optinfo['RATIOxover']=0.5
        if  optinfo['FCxover'] == 'Heuristic':    
            optinfo['RATIOxover']=1.2
        
        
        return optinfo
            
    def __init__(self, worker_index,
                       X0,
                       Xmin = None,
                       Xmax = None,
                       optinfo = None,
                       returninfo = False):
        self.worker_index = worker_index
        self.optinfo = optinfo
    

        if (float(optinfo['_numprocesses'])/optinfo['_group_count'])<=1:
            self.Island_Flag=0
        else:
            self.Island_Flag=1
        
        

            
            

        self.returninfo = returninfo
        if not(self.returninfo):
            self.info = None # contains the information about the run
        else:
            self.info = dict([])
        
        
        
        self.D, self.N = X0.shape   #dimension of the parameter space,  number of individual in population living in this worker
        
        
        self.X = X0
        self.Ntotal=self.optinfo['_group_size']
        self.NbrElite=int(ceil(self.optinfo['PROPelite']*self.N))
        self.Nbrxover=int(floor(self.optinfo['PROPxover']*self.N))
        self.Nbrmutation=self.N-self.NbrElite-self.Nbrxover
        self.Nbroffspring=self.Nbrxover+self.Nbrmutation
        self.NbrMigrants=int(ceil(self.Ntotal*self.optinfo['PCMigration']))
        self.NbrParents=self.N#int(ceil(self.N*self.optinfo['PCParents']))
        # Boundaries
        
        if Xmin is None:
            Xmin = -inf*ones(self.D)
        if Xmax is None:
            Xmax = inf*ones(self.D)
        self.Xmin = Xmin
        self.Xmax = Xmax

        
        if  self.optinfo['FCmutation'] == 'Gaussian':

            if (Xmin!=-inf*ones(self.D)).any() and (Xmax!=inf*ones(self.D)).any():
                self.sigmaMutation=self.optinfo['SCALEmutation']*(Xmax-Xmin)
            else:      # if boudneries are infinite
                self.sigmaMutation=self.optinfo['SCALEmutation']*(10*X0.max(axis=1)-(-10*X0.min(axis=1)))
       
        # set the boundaries for the whole population
        self.Xmin = tile(Xmin.reshape((-1,1)), (1, self.N))   ###Here, shloud be changed to (1, self.Nbroffspring)) 
        self.Xmax = tile(Xmax.reshape((-1,1)), (1, self.N))
        

        
        if self.returninfo:
            self.info['fitness_matrix'] = None
        
    def iterate(self, fitness, global_state):
        
        if global_state is not None:
            self.X_immigrant, migration_flag,self.iter,self.X_bestOverIslands,self.fitness_bestOverIslands = global_state  #shoudld not need globalstate now

        elif self.Island_Flag==0:
            migration_flag=-1
            self.iter=1
        elif self.Island_Flag==1:
            migration_flag=0
            self.iter=1
        ##  sort population from best to worst  ->> in the following we work with sorted X with respect to the fitness
        self.indices_Population_sorted=argsort(fitness)[::-1]
        self.X=self.X[:,self.indices_Population_sorted]   
        self.fitness=fitness[self.indices_Population_sorted]
        
       # print self.X[:,self.NbrElite]
        ##rescale fitness
        self.fitness_scaled=self.scale_fitness()
        
        ### select parents ######
        self.parents_indices=self.select_parents()
        
        #compute the island elite from the fitness that has been just computed for X
        self.fitness_lElite,self.X_lElite,self.X_bestIsland,self.fitness_bestIsland=self.find_elite()
        #print self.X_lElite[:,-1]

        if self.Island_Flag==1:
            ### check migration events
            local_state=self.migration_policy(migration_flag)
        else:
            local_state=-1,[],self.X_bestIsland,self.fitness_bestIsland
            
        ## reproduction: recombine parents with xover and mutation -->compute a new X from the selected adults
        self.X=self.recombine()
       # print self.X[:,-1]
        # Boundary checking
        self.X = maximum(self.X, self.Xmin)
        self.X = minimum(self.X, self.Xmax)
        
 
        return local_state
    
         
    
    @staticmethod
    def combine_local_states(local_states, global_state, optinfo):
        """
        Combines the local states of every worker and returns the global state.
        
        """
        if  global_state is None:      
            iter=1
        else:
            iter=global_state[2]
        
 
        ##find the best fitness over all islands
        fitness=zeros((len(local_states)))
        for i, local_state in zip(range(len(local_states)), local_states):
            fitness[i]=local_state[3]  
        
        indices_bestOverIslands =argmax(fitness)
        fitness_bestOverIslands =fitness[indices_bestOverIslands]
        X_bestOverIslands=local_states[indices_bestOverIslands][2]
        
        migration_flag=local_states[0][0] ## check what is the state of migration

        if migration_flag!=-1:
            #######################################################################################
            
    
            ## update the migration policy
            if migration_flag==0: ## migration hasn't happened 
                 iter=iter+1
                 if iter==optinfo['Minterval']: ## it is time to have a migration
                     migration_flag=1
                     iter=0
                 return ([], migration_flag,iter,X_bestOverIslands,fitness_bestOverIslands)
             
            if migration_flag==1: ## islands have sent their best individuals -> migration will happen
                
                iter=0 #the counter for migration interval is reset
                X_emigrant=[]
                ## local_state=1,self.X_MigrantBest,X_bestIsland,fitness_bestIsland
                for local_state in local_states:
                    X_emigrant.append(local_state[1])
                migration_flag=2 
        
                return (X_emigrant,migration_flag,iter,X_bestOverIslands,fitness_bestOverIslands)
            
            if migration_flag==2:  ## migration has happened 
                iter=iter+1
                migration_flag=0
                return ([], migration_flag,iter,X_bestOverIslands,fitness_bestOverIslands)
            ##########################################################################################
        else:
            return ([], -1,iter,X_bestOverIslands,fitness_bestOverIslands)

         
         
    @staticmethod
    def get_best_fitness(global_state):
        
           
        return global_state[4]
     
    def terminate(self):
        """
        Returns the optimization info if requested, and clears the memory.
        """
        return self.info

    def return_result(self):
        """
        Return the best position and the best fitness.
        """
        return self.X_bestOverIslands, self.fitness_bestOverIslands
    
    def scale_fitness(self):
        ### rescale the fitness values of the entire population###    
        if self.optinfo['FCscale'] == 'Ranking':
            fitness_scaled=arange(self.N,0,-1.)
        return fitness_scaled
    
    
    def select_parents(self):
                #In a strict generational replacement scheme
        #the size of the mating pool is always equal to the size of the population.
        
        ### Stochastic univerasl uniform
        if self.optinfo['FCselection'] == 'StocUnfiform':

            wheel = cumsum(self.fitness_scaled/sum(self.fitness_scaled))  #wheel has the length of the entire population
            parents_indices = zeros((self.NbrParents),'int')
            stepSize = 1./self.NbrParents
            position = rand(1) * stepSize    #  initisialtion position pointer at the starting position
            lowest = 1

            for iparent in range((self.NbrParents)): 
                for ipos in arange(lowest,wheel.shape[0]): # find the wheel position
                    if(position < wheel[ipos]): #if the step fall in this chunk ipos of the wheel
                        parents_indices[iparent] = ipos
                        lowest = ipos
                        break
                position = position + stepSize
                
        return parents_indices 
     
     
    def find_elite(self):
        fitness_lElite= self.fitness[:self.NbrElite]
        X_lElite= self.X[:,:self.NbrElite]
        fitness_bestIsland=self.fitness[0]
        X_bestIsland=self.X[:,0]
        return fitness_lElite,X_lElite,X_bestIsland,fitness_bestIsland
    
    
    def migration_policy(self,migration_flag):   
        ### check migration events
        if migration_flag==0: #migration will not happen at the next step
             local_state=0,[],self.X_bestIsland,self.fitness_bestIsland # nothing is sent
             
        if migration_flag==1:
            # if migration is going to happen at the next step
            #compute the M best from the fitness that has been just computed for X (the emigrant)
            self.fitness_emmigrant= self.fitness[:self.NbrMigrants]
            self.X_emmigrant= self.X[:,:self.NbrMigrants]
            local_state=1,self.X_emmigrant,self.X_bestIsland,self.fitness_bestIsland
            print 'Migration happened'
            
        if migration_flag==2:
        #migration happens, i.e. the immigrant arrive-> change the M worst one with the M best one of worker i-1
            ## choose the M worst
            NbrWorkers=len(self.X_immigrant)/self.N
            self.X[:,-self.NbrMigrants:]=self.X_immigrant[(mod(self.worker_index,NbrWorkers)+1)]  # modulo the number of worker
            local_state=2,[],self.X_bestIsland,self.fitness_bestIsland  # nothing is sent
        return local_state
    
    
    
    def recombine(self):
        Xnew=zeros((self.D,self.N))     
        
        
        Xnew=self.do_crossover(Xnew)
        Xnew=self.do_mutation(Xnew)
        Xnew=self.include_elite(Xnew)
        return Xnew 
    
    def do_crossover(self,Xnew):
        ### CROSSOVER
        if self.optinfo['FCxover']=='DiscreteRandom': 
            for ixover in range(self.Nbrxover):
                parent1_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                parent2_ind=parent1_ind
                while parent1_ind==parent2_ind:     #make sure the two parents are not the same
                    parent2_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                vec=randint(0,2,self.D)
                Xnew[nonzero(vec),ixover]=self.X[nonzero(vec),parent1_ind]
                vec=abs(vec-1)
                Xnew[nonzero(vec),ixover]=self.X[nonzero(vec),parent2_ind]
            
        if self.optinfo['FCxover']=='Onepoint':  
            for ixover in range(self.Nbrxover):
                parent1_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                parent2_ind=parent1_ind
                while parent1_ind==parent2_ind:     #make sure the two parents are not the same
                    parent2_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                    
                split_point=randint(self.D,size=1)[0]
                if split_point != 0:
                    Xnew[:split_point,ixover]=self.X[:split_point,parent1_ind]
                else :
                    Xnew[split_point,ixover]=self.X[split_point,parent1_ind]
                if split_point != self.D:
                    Xnew[split_point:,ixover]=self.X[split_point:,parent2_ind]    
                else :
                    Xnew[split_point,ixover]=self.X[split_point,parent2_ind]    
                
        if self.optinfo['FCxover']=='TwoPoints': 
            for ixover in range(self.Nbrxover):
                parent1_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                parent2_ind=parent1_ind
                while parent1_ind==parent2_ind:     #make sure the two parents are not the same
                    parent2_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]

                split_points1=randint(self.D,size=1)
                split_points2=split_points1
                while split_points1==split_points2:
                    split_points2=randint(self.D,size=1)
                split_points=sort(array([split_points1,split_points2])  )

                if split_points[0] != 0:
                    Xnew[:split_points[0],ixover]=self.X[:split_points[0],parent1_ind]
                else :
                    Xnew[split_points[0],ixover]=self.X[split_points[0],parent1_ind]
                    
                Xnew[split_points[0]:split_points[1],ixover]=self.X[split_points[0]:split_points[1],parent2_ind] 
                
                if split_points[1] != self.D:
                    Xnew[split_points[1]:,ixover]=self.X[split_points[1]:,parent1_ind]
                else :
                    Xnew[split_points[1],ixover]=self.X[split_points[1],parent1_ind]
                         
        if self.optinfo['FCxover']=='Heuristic': 
            for ixover in range(self.Nbrxover):
                parent1_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                parent2_ind=parent1_ind
                while parent1_ind==parent2_ind:     #make sure the two parents are not the same
                    parent2_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                
                if self.fitness[parent1_ind]>=self.fitness[parent2_ind]: # choose the best parent
                    Xnew[:,ixover]=self.X[:,parent2_ind]+self.optinfo['RATIOxover']*(self.X[:,parent1_ind]-self.X[:,parent2_ind])
                else:
                    Xnew[:,ixover]=self.X[:,parent1_ind]+self.optinfo['RATIOxover']*(self.X[:,parent2_ind]-self.X[:,parent1_ind])
                    
        if self.optinfo['FCxover']=='Intermediate': 
            for ixover in range(self.Nbrxover):
                parent1_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                parent2_ind=parent1_ind
                while parent1_ind==parent2_ind:     #make sure the two parents are not the same
                    parent2_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                Xnew[:,ixover]=self.X[:,parent1_ind]+(-0.25+1.25*rand(self.D))*(self.X[:,parent2_ind]-self.X[:,parent1_ind])
              
        if self.optinfo['FCxover']=='LinearComb':
            for ixover in range(self.Nbrxover):
                parent1_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                parent2_ind=parent1_ind
                while parent1_ind==parent2_ind:     #make sure the two parents are not the same
                    parent2_ind=self.parents_indices[randint(self.NbrParents,size=1)[0]]
                    
                Xnew[:,ixover]=self.X[:,parent1_ind]+self.optinfo['RATIOxover']*(self.X[:,parent2_ind]-self.X[:,parent1_ind])
        return Xnew
    
    def do_mutation(self,Xnew):
                #### MUTATION
        if self.optinfo['FCmutation']=='Gaussian':
            for imut in range(self.Nbrmutation):
                Xnew[:,self.Nbrxover+imut]=self.X[:,self.parents_indices[randint(self.NbrParents,size=1)[0]]]+self.sigmaMutation*randn(self.D)
            self.sigmaMutation=self.sigmaMutation*(1-self.optinfo['shrinkMutation']*self.iter/self.optinfo['iterations'])
                 
        if self.optinfo['FCmutation']=='Uniform':
            
            for imut in range(self.Nbrmutation):
                Xnew[:,self.Nbrxover+imut]=self.X[:,self.parents_indices[randint(self.NbrParents,size=1)[0]]]+(self.Xmax[:,0]-self.Xmin[:,0])*rand(self.D)
        return Xnew
    
    def include_elite(self,Xnew):
        ### add the current elite to the next  generation

        Xnew[:,self.Nbrxover+self.Nbrmutation:]=self.X_lElite      
        return Xnew
    
############################################################################### 